// MysqlChatMemory.java
package com.zry.ai.utils;


import com.zry.ai.entity.ChatMemoryEntity.ChatMessageEntity;
import com.zry.ai.mapper.MessageMapper;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Component
@RequiredArgsConstructor
public class MysqlChatMemory implements ChatMemory {
    private final MessageMapper messageMapper;
    @Override
    public void add(String conversationId, List<Message> messages) {
        List<ChatMessageEntity> entities = messages.stream()
            .map(msg -> {
                ChatMessageEntity entity = new ChatMessageEntity();
                entity.setConversationId(conversationId);
                //TODO
                entity.setContent(msg.getText());
                
                if (msg instanceof UserMessage) {
                    entity.setMessageType(ChatMessageEntity.MessageType.USER);
                } else if (msg instanceof AssistantMessage) {
                    entity.setMessageType(ChatMessageEntity.MessageType.ASSISTANT);
                }
                
                entity.setCreatedAt(LocalDateTime.now());
                return entity;
            })
            .collect(Collectors.toList());
        
        messageMapper.insertMessages(entities);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<ChatMessageEntity> entities = messageMapper.findLastNMessages(conversationId, lastN);
        Collections.reverse(entities);
        
        return entities.stream()
            .map(entity -> {
                switch (entity.getMessageType()) {
                    case USER:
                        return new UserMessage(entity.getContent());
                    case ASSISTANT:
                        return new AssistantMessage(entity.getContent());
                    default:
                        throw new IllegalArgumentException("未知的消息类型");
                }
            })
            .collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        messageMapper.deleteByConversationId(conversationId);
    }
}